fix: [spearbit-98][quantstamp-8] State is not cached properly before validation/execution steps, Non-Idempotent Pre-Execution Hook #58

adam-alchemy · 2024-01-10T21:39:38Z

Motivation

https://github.com/spearbit-audits/alchemy-nov-review/issues/98

We are inconsistent in what account state we cache into memory prior to execution, and what we load from storage when needed.

Solution

Define distinct validation and execution phases, for which storage changes to account state should not result in different execution.

Update the implementation of UpgradeableModularAccount to cache according to these phases.

This is a very large re-organization of hook logic, because we previously mixed loading from storage and memory. This change isolates all storage loads first, then performs memory loads for executions as each one is needed.

Test the behavior of account state changes to ensure phases are preserved. This is a very tricky to orchestrate test since it involves re-entering the account from different plugin functions, and making assertions about the overall call flow. It also has a very large number of cases to test (89), so please read the comments in the AccountStatePhasesTest carefully.

Implementation strategy for hook caching

(Note the the extra slots for post-only hooks are only allocated if needed, otherwise they are not added to the length).

Regular (non-`IPluginExecutor`) call flow:

Load all pre exec hooks into memory
Allocate a bytes[] preExecReturnData of preExecHooksLength + 1 for holding the return data. (last element goes unused, to indicate no parameter)
Allocate a FunctionReference[][] postExecHooks of preExecHooksLength + 1 for holding the post-exec hooks.
Fill the last sub-array in postExecHooks with the post-only hooks.
Fill in each other sub-array with all associated post-exec hooks.
Run the pre exec hooks and put the values into preExecReturnData
Run the exec function
Run all postExecHooks, passing their return data from the correct array element in preExecReturnData.

Permitted call flow:

Load all pre permitted call hooks and pre exec hooks into memory
Allocate a bytes[] preHookReturnData of preExecHooksLength + prePermittedCallHooksLength + 2 for holding the return data. (Last 2 elements go unused)
Allocate a FunctionReference[][] postHooks of preExecHooksLength + prePermittedCallHooksLength + 2 for holding the post-* hooks.
Fill in the last sub-array with the post-only permitted call hooks.
Fill in the second to last sub-array with the post-only exec hooks.
Fill in each other sub-array with the associated post-permitted call hooks first, then the associated post-exec hooks.
Run the pre permitted call hooks and store the return data.
Run the pre exec hooks and store the return data.
Run the exec function.
Run all postExecHooks, passing their return data from the correct array element in preHookReturnData.

Code size

Code size is now 24464 bytes, leaving 112 bytes to spare.

src/account/UpgradeableModularAccount.sol

jaypaik

I think this lgtm, have to step out for a bit but will take another pass soon!

src/account/UpgradeableModularAccount.sol

jaypaik · 2024-01-12T22:31:30Z

src/account/UpgradeableModularAccount.sol

        bytes4 selector = _selectorFromCallData(userOp.callData);
        SelectorData storage selectorData = _getAccountStorage().selectorData[selector];

        FunctionReference userOpValidationFunction = selectorData.userOpValidation;
-        hasPreValidationHooks = selectorData.hasPreUserOpValidationHooks;
+        bool hasPreValidationHooks = selectorData.hasPreUserOpValidationHooks;


src/account/UpgradeableModularAccount.sol

jaypaik · 2024-01-12T22:44:36Z

src/account/UpgradeableModularAccount.sol

+        postHooksToRun = new FunctionReference[][](postHooksToRunLength);
+        postHookArgs = new bytes[](postHooksToRunLength);
+
+        uint256 currentIndex = 0;


Suggested change

uint256 currentIndex = 0;

uint256 currentIndex;

jaypaik · 2024-01-12T22:49:49Z

src/account/UpgradeableModularAccount.sol

+        postHooksToRun = new FunctionReference[][](postHooksToRunLength);
+        postHookArgs = new bytes[](postHooksToRunLength);
+
+        uint256 currentIndex = 0;


Suggested change

uint256 currentIndex = 0;

uint256 currentIndex;

I want to keep these explicit for the reader, rather than having them thinking about the default values.

jaypaik · 2024-01-12T22:55:10Z

src/account/UpgradeableModularAccount.sol

+        // We use `currentIndex + prePermittedCallHooks.length` for the starting index but do not update it,
+        // because its current value is useful for executing the hooks.
+        _cacheAssociatedPostHooks(
+            preExecHooks, selectorData.executionHooks, postHooksToRun, currentIndex + prePermittedCallHooks.length


Nit: maybe store currentIndex + prePermittedCallHooks.length in a variable to reuse in line 694?

src/account/UpgradeableModularAccount.sol

fangting-alchemy · 2024-01-12T22:58:21Z

src/account/UpgradeableModularAccount.sol

+        // in the args for their empty `bytes` argument.
+        uint256 postHooksToRunLength = preExecHooks.length + (hasPostOnlyExecHooks ? 1 : 0);
+        postHooksToRun = new FunctionReference[][](postHooksToRunLength);
+        postHookArgs = new bytes[](postHooksToRunLength);


Why not keep same indexing for postHooksToRun, postHookArgs, and preHooks, add postOnlyHooks in the end of postHooksToRun?

Save you the need of passing currentIndex.

I guess the indexing wont match when the permitted call hooks are present. With such a change to hooks run rn, I am now more open to remove permitted call hooks.

src/account/UpgradeableModularAccount.sol

fangting-alchemy · 2024-01-12T23:12:34Z

src/account/UpgradeableModularAccount.sol

+        // in the args for their empty `bytes` argument.
+        uint256 postHooksToRunLength = preExecHooks.length + (hasPostOnlyExecHooks ? 1 : 0);
+        postHooksToRun = new FunctionReference[][](postHooksToRunLength);
+        postHookArgs = new bytes[](postHooksToRunLength);


I guess the indexing wont match when the permitted call hooks are present. With such a change to hooks run rn, I am now more open to remove permitted call hooks.

jaypaik

Thanks for the thorough tests! Liking the consolidation of post hooks and post only hooks, along with the design of using a 2d array to address duplicate pre hooks that are associated with different post hooks.

Left some comments, mostly nits.

src/account/UpgradeableModularAccount.sol

test/account/phases/AccountStatePhases.t.sol

src/account/UpgradeableModularAccount.sol

test/mocks/plugins/AccountStateMutatingPlugin.sol

jaypaik · 2024-01-14T03:28:58Z

test/account/phases/AccountStatePhasesExec.t.sol

+    // Source: pre-Exec
+    // Target: pre-UserOp-Validation
+    // n/a - runs before
+
+    // Source: pre-Exec
+    // Target: UserOp-Validation
+    // n/a - runs before
+
+    // Source: pre-Exec
+    // Target: pre-Runtime-Validation
+    // n/a - runs before
+
+    // Source: pre-Exec
+    // Target: Runtime-Validation
+    // n/a - runs before
+
+    // Source: pre-Exec
+    // Target: pre-Exec (same phase)
+    // Addition (first element): *impossible*


Nit: wasn't really clear why these were repeated (here and below) until I realized you're using these as placeholders for tests you are purposely skipping.

rholterhus · 2024-01-15T07:56:00Z

src/account/UpgradeableModularAccount.sol

+        if (hasPostOnlyPermittedCallHooks) {
+            // If we have post-only hooks, we allocate an single FunctionReference[] for them, and one element in
+            // the args for their empty `bytes` argument. We put this into the first element of the post hooks in
+            // order to have it run last.
+            postHooksToRun[currentIndex] =
+                CastLib.toFunctionReferenceArray(permittedCallData.permittedCallHooks.postOnlyHooks.getAll());
            unchecked {
-                maxPostHooksToRunLength += preHookSet.getCount(CastLib.toSetValue(preExecHooks[i]));
-                ++i;
+                ++currentIndex;
+            }
+        }
+
+        if (hasPostOnlyExecHooks) {
+            // If we have post-only hooks, we allocate an single FunctionReference[] for them, and one element in
+            // the args for their empty `bytes` argument. We put this into the first element of the post hooks in
+            // order to have it run last.
+            postHooksToRun[currentIndex] =
+                CastLib.toFunctionReferenceArray(selectorData.executionHooks.postOnlyHooks.getAll());
+            unchecked {
+                ++currentIndex;
            }
        }


This logic will run the different hooks in this order:

Associated post hooks for selectorData.executionHooks

Associated post hooks for permittedCallData.permittedCallHooks

Post only hooks for selectorData.executionHooks

Post only hooks for permittedCallData.permittedCallHooks

Technically, I think the order before was 1, 3, 2, 4. For example see this code from before:

// Run any post exec hooks for the `executeFromPluginExternal` selector _doCachedPostHooks(postExecHooks, postExecHookArgs); // Run any post only exec hooks for the `executeFromPluginExternal` selector _doCachedPostHooks( CastLib.toFunctionReferenceArray( storage_.selectorData[IPluginExecutor.executeFromPluginExternal.selector] .executionHooks .postOnlyHooks .getAll() ), new bytes[](0) ); // Run any post permitted call hooks specific to this caller and the `executeFromPluginExternal` selector _doCachedPostHooks(postPermittedCallHooks, postPermittedCallHookArgs); // Run any post only permitted call hooks specific to this caller and the `executeFromPluginExternal` // selector _doCachedPostHooks( CastLib.toFunctionReferenceArray( storage_.permittedCalls[permittedCallKey].permittedCallHooks.postOnlyHooks.getAll() ), new bytes[](0) );

I think either way is okay, but is this an intentional change?

Yeah, this is probably worth fixing. We'll need to keep track of the two indices that point to the beginning of each type of associated hook but seems like an easy fix.

I fixed this to run in the post hooks in the following order:

associated postExec

post-only postExec

associated postPermittedCall

post-only postPermittedCall

In this commit: bae9391

Could you please double check? I walked through with a tracer but wanted to confirm.

It seems that permitted call hooks are being removed (basing this off of PR 65). Will this PR need to be updated as well (maybe after merging PR 65)?

PR 65 will be rebased on top of this and merged later. So I guess we don't strictly need to review this, since the invocation of _doPrePermittedCallAndPreExecHooks will be replaced with _doPreExecHooks.

New PR addressing this: #76 (#65 has been closed).

…validation/execution steps, Non-Idempotent Pre-Execution Hook (#58)

adam-alchemy requested a review from fangting-alchemy January 10, 2024 21:39

adamegyed added 15 commits January 10, 2024 18:29

Test development progress

3ccd7aa

Test setup progress

2efeb4b

All user op validation state changing tests

fcdf9c6

Runtime validation tests complete

ce8b5e6

partial pre exec hook tests

35f3fab

Update tests with new revert message

3a1db4d

Add all pre exec state change tests

68df750

cache account state phases for exec hooks

1bde18c

Caching implementation

f2c205c

Simplify _doPreExecHooks to reduce code size

c6bd94c

pre hook execution optimization

a26c975

Complete ASP test cases

012ebd5

Split up account phases tests

017c9a6

Docs update and refactor

a7ba56e

Add comment

ac8c8c0

adam-alchemy force-pushed the adam/fix-98 branch from ffabbae to ac8c8c0 Compare January 10, 2024 23:30

gpersoon reviewed Jan 11, 2024

View reviewed changes

src/account/UpgradeableModularAccount.sol Outdated Show resolved Hide resolved

adamegyed added 2 commits January 11, 2024 11:33

Correct usage of index in _doPreExecHooks

37104e8

lint: fix function order

e15c8ef

howydev reviewed Jan 12, 2024

View reviewed changes

src/account/UpgradeableModularAccount.sol Show resolved Hide resolved

jaypaik reviewed Jan 12, 2024

View reviewed changes

fangting-alchemy reviewed Jan 12, 2024

View reviewed changes

jaypaik approved these changes Jan 14, 2024

View reviewed changes

rholterhus reviewed Jan 15, 2024

View reviewed changes

adamegyed added 2 commits January 16, 2024 12:38

Comment fixes

4f4263b

post hook ordering fix

bae9391

fangting-alchemy mentioned this pull request Jan 18, 2024

feat: remove permitted call hooks from MSCA (rebase on cache PR) #76

Merged

adam-alchemy merged commit 28c7bf1 into audit-2023-11-20 Jan 18, 2024
3 checks passed

adam-alchemy deleted the adam/fix-98 branch January 18, 2024 16:36

adam-alchemy added a commit that referenced this pull request Jan 19, 2024

fix: [spearbit-98][quantstamp-8] State is not cached properly before …

1942415

…validation/execution steps, Non-Idempotent Pre-Execution Hook (#58)

jaypaik pushed a commit that referenced this pull request Jan 25, 2024

fix: [spearbit-98][quantstamp-8] State is not cached properly before …

90337d6

…validation/execution steps, Non-Idempotent Pre-Execution Hook (#58)

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

fix: [spearbit-98][quantstamp-8] State is not cached properly before validation/execution steps, Non-Idempotent Pre-Execution Hook #58

fix: [spearbit-98][quantstamp-8] State is not cached properly before validation/execution steps, Non-Idempotent Pre-Execution Hook #58

adam-alchemy commented Jan 10, 2024 •

edited

Loading

jaypaik left a comment

jaypaik Jan 12, 2024

jaypaik Jan 12, 2024

jaypaik Jan 12, 2024

adam-alchemy Jan 16, 2024

jaypaik Jan 12, 2024 •

edited

Loading

fangting-alchemy Jan 12, 2024

fangting-alchemy Jan 12, 2024

fangting-alchemy Jan 12, 2024

jaypaik left a comment

jaypaik Jan 14, 2024

rholterhus Jan 15, 2024

jaypaik Jan 16, 2024

adam-alchemy Jan 16, 2024

rholterhus Jan 17, 2024

adam-alchemy Jan 17, 2024

jaypaik Jan 21, 2024

fix: [spearbit-98][quantstamp-8] State is not cached properly before validation/execution steps, Non-Idempotent Pre-Execution Hook #58

fix: [spearbit-98][quantstamp-8] State is not cached properly before validation/execution steps, Non-Idempotent Pre-Execution Hook #58

Conversation

adam-alchemy commented Jan 10, 2024 • edited Loading

Motivation

Solution

Implementation strategy for hook caching

Regular (non-IPluginExecutor) call flow:

Permitted call flow:

Code size

jaypaik left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jaypaik Jan 12, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jaypaik left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

adam-alchemy commented Jan 10, 2024 •

edited

Loading

Regular (non-`IPluginExecutor`) call flow:

jaypaik Jan 12, 2024 •

edited

Loading